package com.leilei;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.apache.flink.api.common.RuntimeExecutionMode;
import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.windowing.assigners.TumblingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.util.Collector;

import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

/**
 * @author lei
 * @version 1.0
 * @date 2021/4/11 14:05
 * @desc 数据关联 ，关联不上的数据仍会保留下来
 */
public class CoGroupOperator {
    public static void main(String[] args) throws Exception {
        List<User> users = Arrays.asList(new User("张三", 12),
                new User("李四", 13), new User("王五", 14));

        List<Student> students = Arrays.asList(new Student("张三", "小三"),
                new Student("李四", "小四"),new Student("赵六", "小六子"));

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.setRuntimeMode(RuntimeExecutionMode.AUTOMATIC);
        DataStreamSource<User> userStream = env.fromCollection(users);
        DataStreamSource<Student> studentStream = env.fromCollection(students);
        // 对两个数据源进行关联
        DataStream<Tuple3<String, String, Integer>> result = userStream.coGroup(studentStream)
                // 关联条件 user中 name 要等于 student中name
                .where(u -> u.name).equalTo(t -> t.name)
                // 这里示意是使用时间滚动窗口 1秒计算一次
                .window(TumblingProcessingTimeWindows.of(Time.seconds(1)))
                // 数据处理  需注意，并非where满足才会执行，where不满足（未关联上）也会调用，只是某一只迭代器无数据罢了，需要自己处理
                .apply(new CoGroupFunction<User, Student, Tuple3<String, String, Integer>>() {
                    @Override
                    public void coGroup(Iterable<User> first, Iterable<Student> second, Collector<Tuple3<String, String, Integer>> out) throws Exception {
                        Iterator<User> users = first.iterator();
                        Iterator<Student> students = second.iterator();
                        if (users.hasNext() && students.hasNext()) {
                                User user = users.next();
                                Student student = students.next();
                                out.collect(Tuple3.of(user.name, student.nickName, user.age));
                        } else if (users.hasNext() || students.hasNext()) {
                            if (users.hasNext()) {
                                while (users.hasNext()) {
                                    User user = users.next();
                                    out.collect(Tuple3.of(user.name, "无昵称", user.age));
                                }
                            }
                            if (students.hasNext()) {
                                while (students.hasNext()) {
                                    Student student = students.next();
                                    out.collect(Tuple3.of(student.name, student.nickName, -1));
                                }
                            }

                        }
                    }
                });
        result.print("coGroup");
        env.execute();
    }

    @NoArgsConstructor
    @AllArgsConstructor
    @Data
    @Builder
    public static class User {
        private String name;
        private Integer age;
    }

    @NoArgsConstructor
    @AllArgsConstructor
    @Data
    @Builder
    public static class Student {
        private String name;
        private String nickName;
    }
}
